#include "MeshConverter.h"
#include "resource.h"
#include <stdio.h>

D3DVERTEXELEMENT9 basicElements[] = 
{
	{ 0, 0,  D3DDECLTYPE_FLOAT3, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_POSITION, 0 },
	{ 0, 24, D3DDECLTYPE_FLOAT2, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_TEXCOORD, 0 },
	D3DDECL_END()
};

D3DVERTEXELEMENT9 normalElements[] = 
{
	{ 0, 0,  D3DDECLTYPE_FLOAT3, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_POSITION, 0 },
	{ 0, 12, D3DDECLTYPE_FLOAT3, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_NORMAL,   0 },
	{ 0, 24, D3DDECLTYPE_FLOAT2, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_TEXCOORD, 0 },
	D3DDECL_END()
};

D3DVERTEXELEMENT9 fullElements[] = 
{
	{ 0, 0,  D3DDECLTYPE_FLOAT3, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_POSITION, 0 },
	{ 0, 12, D3DDECLTYPE_FLOAT3, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_NORMAL,   0 },
	{ 0, 24, D3DDECLTYPE_FLOAT2, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_TEXCOORD, 0 },
	{ 0, 32, D3DDECLTYPE_FLOAT3, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_BINORMAL,  0 },
	{ 0, 44, D3DDECLTYPE_FLOAT3, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_TANGENT, 0 },
	D3DDECL_END()
};

LRESULT CALLBACK WndProc(HWND hWnd, UINT message, WPARAM wParam, LPARAM lParam)
{
	switch (message)
	{
		case WM_DESTROY:
			PostQuitMessage(0);
			break;
		default:
			return DefWindowProc(hWnd, message, wParam, lParam);
	}

	return 0;
}

HRESULT InitializeWindow(HWND* pHwnd)
{
	HINSTANCE hInstance = GetModuleHandle(NULL);

	WNDCLASSEX wcex;
	wcex.cbSize = sizeof(WNDCLASSEX);
	wcex.style = CS_HREDRAW | CS_VREDRAW;
	wcex.lpfnWndProc = WndProc;
	wcex.cbClsExtra = 0;
	wcex.cbWndExtra = 0;
	wcex.hInstance = hInstance;
	wcex.hIcon = LoadIcon(hInstance, (LPCTSTR)IDI_ICON1);
	wcex.hIconSm = LoadIcon(hInstance, (LPCTSTR)IDI_ICON1);
	wcex.hCursor = LoadCursor(NULL, IDC_ARROW);
	wcex.hbrBackground = (HBRUSH)(COLOR_WINDOW+1);
	wcex.lpszMenuName = NULL;
	wcex.lpszClassName = L"D3D10Engine";

	if (!RegisterClassEx(&wcex)) return E_FAIL;

	RECT rc = {0, 0, 640, 480};
	AdjustWindowRect(&rc, WS_OVERLAPPEDWINDOW, FALSE);
	*pHwnd = CreateWindow(L"D3D10Engine", L"D3D10Engine", WS_OVERLAPPEDWINDOW, CW_USEDEFAULT, CW_USEDEFAULT, rc.right-rc.left, rc.bottom - rc.top, NULL, NULL, hInstance, NULL);

	if (!*pHwnd) return E_FAIL;

	return S_OK;
}

HRESULT ReadOptions(int argc, LPWSTR args[], Options* pOptions)
{
	for (int currentArg = 0; currentArg < argc; currentArg++)
	{
		if (wcscmp(args[currentArg], L"-iF") == 0)
		{
			currentArg++;
			if (currentArg < argc)
			{
				int len = wcslen(args[currentArg]) + 1;
				pOptions->inputPath = (LPWSTR)malloc(sizeof(WCHAR) * len);
				wcscpy_s(pOptions->inputPath, len, args[currentArg]);
			}
			else
			{
				return E_FAIL;
			}
		}
		else if (wcscmp(args[currentArg], L"-oF") == 0)
		{
			currentArg++;
			if (currentArg < argc)
			{
				int len = wcslen(args[currentArg]) + 1;
				pOptions->outputPath = (LPWSTR)malloc(sizeof(WCHAR) * len);
				wcscpy_s(pOptions->outputPath, len, args[currentArg]);
			}
			else
			{
				return E_FAIL;
			}
		}
		else if (wcscmp(args[currentArg], L"-nT") == 0)
		{
			currentArg++;
			if (currentArg < argc)
			{
				int len = wcslen(args[currentArg]) + 1;
				pOptions->normalTexture = (LPWSTR)malloc(sizeof(WCHAR) * len);
				wcscpy_s(pOptions->normalTexture, len, args[currentArg]);
			}
			else
			{
				return E_FAIL;
			}
		}
		else if (wcscmp(args[currentArg], L"-sT") == 0)
		{
			currentArg++;
			if (currentArg < argc)
			{
				int len = wcslen(args[currentArg]) + 1;
				pOptions->specularTexture = (LPWSTR)malloc(sizeof(WCHAR) * len);
				wcscpy_s(pOptions->specularTexture, len, args[currentArg]);
			}
			else
			{
				return E_FAIL;
			}
		}
		else if (wcscmp(args[currentArg], L"-gN") == 0)
		{
			pOptions->generateNormals = true;
		}
		else if (wcscmp(args[currentArg], L"-gNTB") == 0)
		{
			pOptions->generateNormalTangentBinormals = true;
		}
	}

	return S_OK;
}

HRESULT Initialize(Options options, LPDIRECT3DDEVICE9* ppDevice)
{
	UNREFERENCED_PARAMETER(options);

	HRESULT hr = S_OK;
	HWND hwnd;
	LPDIRECT3D9 pD3D9 = NULL;
	D3DPRESENT_PARAMETERS pParams;
	
	hr = InitializeWindow(&hwnd);
	if (FAILED(hr)) goto END;

	pD3D9 = Direct3DCreate9(D3D_SDK_VERSION);
	if (pD3D9 == NULL)
	{
		hr = E_FAIL;
		goto END;
	}

	ZeroMemory(&pParams, sizeof(D3DPRESENT_PARAMETERS));
	pParams.Windowed = true;
	pParams.hDeviceWindow = hwnd;
	pParams.SwapEffect = D3DSWAPEFFECT_DISCARD;

	hr = pD3D9->CreateDevice(D3DADAPTER_DEFAULT, D3DDEVTYPE_HAL, hwnd, D3DCREATE_SOFTWARE_VERTEXPROCESSING, &pParams, ppDevice);
	if (FAILED(hr)) goto END;

END:
	if (pD3D9 != NULL) pD3D9->Release();

	if (FAILED(hr)) return hr;
	return S_OK;
}

HRESULT Load(LPDIRECT3DDEVICE9 pDevice, Options options, LPD3DXMATERIAL* pMaterials, DWORD* nMaterials, LPD3DXMESH* pMesh)
{
	HRESULT hr = S_OK;
	LPD3DXBUFFER adjacencyBuffer = NULL;
	LPD3DXBUFFER materialBuffer = NULL;
	LPD3DXBUFFER effectBuffer = NULL;
	LPD3DXMATERIAL materialsCopy;
	LPD3DXMESH tempMesh = NULL;
	LPD3DXMESH copyMesh = NULL;

	if (options.inputPath == NULL)
	{
		hr = E_FAIL;
		goto END;
	}

	hr = D3DXLoadMeshFromX(options.inputPath, NULL, pDevice, &adjacencyBuffer, &materialBuffer, &effectBuffer, nMaterials, &tempMesh);
	if (FAILED(hr)) goto END;

	*pMaterials = (LPD3DXMATERIAL)malloc(sizeof(D3DXMATERIAL) * (*nMaterials));
	ZeroMemory(*pMaterials, sizeof(D3DXMATERIAL) * (*nMaterials));
	materialsCopy = (LPD3DXMATERIAL)materialBuffer->GetBufferPointer();
	memcpy(*pMaterials, materialsCopy, sizeof(D3DXMATERIAL) * (*nMaterials));
	for (UINT i = 0; i < *nMaterials; i++)
	{
		if (materialsCopy[i].pTextureFilename != NULL)
		{
			int len = strlen(materialsCopy[i].pTextureFilename) + 1;
			(*pMaterials)[i].pTextureFilename = (LPSTR)malloc(sizeof(char) * len);
			strcpy_s((*pMaterials)[i].pTextureFilename, len, materialsCopy[i].pTextureFilename);
		}
	}

	D3DVERTEXELEMENT9* vertDecl;
	if (options.generateNormalTangentBinormals)
	{
		vertDecl = fullElements;
	}
	else if (options.generateNormals)
	{
		vertDecl = normalElements;
	}
	else
	{
		vertDecl = basicElements;
	}
	
	if (options.generateNormalTangentBinormals)
	{
		hr = tempMesh->CloneMesh(NULL, vertDecl, pDevice, &copyMesh);
		tempMesh->Release();
		tempMesh = NULL;
		if (FAILED(hr)) goto END;
		hr = D3DXComputeTangentFrameEx(copyMesh,
		                               D3DDECLUSAGE_TEXCOORD, 0,
			    					   D3DDECLUSAGE_BINORMAL, 0,
				    				   D3DDECLUSAGE_TANGENT, 0,
					    			   D3DDECLUSAGE_NORMAL, 0,
						    		   D3DXTANGENT_CALCULATE_NORMALS,
							    	   NULL,
								       .01f, .25f, .02f,
								       pMesh, NULL);
		if (FAILED(hr)) goto END;
	}
	else
	{
		hr = tempMesh->CloneMesh(NULL, vertDecl, pDevice, pMesh);
	}

END:
	if (adjacencyBuffer != NULL) adjacencyBuffer->Release();
	if (materialBuffer != NULL) materialBuffer->Release();
	if (effectBuffer != NULL) effectBuffer->Release();
	if (tempMesh != NULL) tempMesh->Release();
	if (copyMesh != NULL) copyMesh->Release();

	if (FAILED(hr)) return hr;
	return S_OK;	
}

HRESULT Convert(LPD3DXMATERIAL materials, DWORD nMaterials, LPD3DXMESH mesh, Options options, SDKMESH* pSdkMesh)
{
	HRESULT hr = S_OK;
	LPDIRECT3DINDEXBUFFER9 indexBuffer = NULL;
	LPDIRECT3DVERTEXBUFFER9 vertexBuffer = NULL;
	D3DVERTEXELEMENT9 declaration[MAX_FVF_DECL_SIZE];
	D3DVERTEXELEMENT9 end = D3DDECL_END();
	
	pSdkMesh->meshHeader = (SDKMESH_HEADER*)malloc(sizeof(SDKMESH_HEADER));
	pSdkMesh->vertexHeaders = (SDKMESH_VERTEX_BUFFER_HEADER*)malloc(sizeof(SDKMESH_VERTEX_BUFFER_HEADER));
	pSdkMesh->indexHeaders = (SDKMESH_INDEX_BUFFER_HEADER*)malloc(sizeof(SDKMESH_INDEX_BUFFER_HEADER));
	pSdkMesh->materials = (SDKMESH_MATERIAL*)malloc(sizeof(SDKMESH_MATERIAL)*nMaterials);
	pSdkMesh->subsets = (SDKMESH_SUBSET*)malloc(sizeof(SDKMESH_SUBSET));
	pSdkMesh->meshes = (SDKMESH_MESH*)malloc(sizeof(SDKMESH_MESH));

	ZeroMemory(pSdkMesh->meshHeader, sizeof(SDKMESH_HEADER));
	ZeroMemory(pSdkMesh->vertexHeaders, sizeof(SDKMESH_VERTEX_BUFFER_HEADER));
	ZeroMemory(pSdkMesh->indexHeaders, sizeof(SDKMESH_INDEX_BUFFER_HEADER));
	ZeroMemory(pSdkMesh->materials, sizeof(SDKMESH_MATERIAL) * nMaterials);
	ZeroMemory(pSdkMesh->subsets, sizeof(SDKMESH_SUBSET));
	ZeroMemory(pSdkMesh->meshes, sizeof(SDKMESH_MESH));

	for (UINT i = 0; i < nMaterials; i++)
	{
		pSdkMesh->materials[i].Ambient.x = materials[i].MatD3D.Ambient.r;
		pSdkMesh->materials[i].Ambient.y = materials[i].MatD3D.Ambient.g;
		pSdkMesh->materials[i].Ambient.z = materials[i].MatD3D.Ambient.b;
		pSdkMesh->materials[i].Ambient.w = materials[i].MatD3D.Ambient.a;

		pSdkMesh->materials[i].Diffuse.x = materials[i].MatD3D.Diffuse.r;
		pSdkMesh->materials[i].Diffuse.y = materials[i].MatD3D.Diffuse.g;
		pSdkMesh->materials[i].Diffuse.z = materials[i].MatD3D.Diffuse.b;
		pSdkMesh->materials[i].Diffuse.w = materials[i].MatD3D.Diffuse.a;

		pSdkMesh->materials[i].Specular.x = materials[i].MatD3D.Specular.r;
		pSdkMesh->materials[i].Specular.y = materials[i].MatD3D.Specular.g;
		pSdkMesh->materials[i].Specular.z = materials[i].MatD3D.Specular.b;
		pSdkMesh->materials[i].Specular.w = materials[i].MatD3D.Specular.a;

		pSdkMesh->materials[i].Emissive.x = materials[i].MatD3D.Emissive.r;
		pSdkMesh->materials[i].Emissive.y = materials[i].MatD3D.Emissive.g;
		pSdkMesh->materials[i].Emissive.z = materials[i].MatD3D.Emissive.b;
		pSdkMesh->materials[i].Emissive.w = materials[i].MatD3D.Emissive.a;

		if (materials[i].pTextureFilename != 0) strcpy_s(pSdkMesh->materials[i].DiffuseTexture, strlen(materials[i].pTextureFilename) + 1, materials[i].pTextureFilename);
		if (options.normalTexture != NULL)
		{
			size_t i = 0;
			wcstombs_s(&i, pSdkMesh->materials[i].NormalTexture, MAX_TEXTURE_NAME, options.normalTexture, wcslen(options.normalTexture) + 1);
		}
		if (options.specularTexture != NULL)
		{
			size_t i = 0;
			wcstombs_s(&i, pSdkMesh->materials[i].SpecularTexture, MAX_TEXTURE_NAME, options.specularTexture, wcslen(options.specularTexture) + 1);
		}

		pSdkMesh->materials[i].Power = materials[i].MatD3D.Power;
	}

	hr = mesh->GetIndexBuffer(&indexBuffer);
	if (FAILED(hr)) goto END;
	D3DINDEXBUFFER_DESC iDesc;
	hr = indexBuffer->GetDesc(&iDesc);
	if (FAILED(hr)) goto END;
	pSdkMesh->indexHeaders[0].IndexType = iDesc.Type;
	pSdkMesh->indexHeaders[0].SizeBytes = iDesc.Size;
	pSdkMesh->indexHeaders[0].NumIndices = mesh->GetNumFaces() * 3;

	hr = mesh->GetVertexBuffer(&vertexBuffer);
	if (FAILED(hr)) goto END;
	D3DVERTEXBUFFER_DESC vDesc;
	hr = vertexBuffer->GetDesc(&vDesc);
	if (FAILED(hr)) goto END;
	hr = mesh->GetDeclaration(declaration);
	if (FAILED(hr)) goto END;
	for (UINT i = 0; memcmp(&declaration[i], &end, sizeof(D3DVERTEXELEMENT9)) != 0 ; i++)
	{
		memcpy(&pSdkMesh->vertexHeaders[0].d3dVertexElements9[i], &declaration[i], sizeof(D3DVERTEXELEMENT9));
	}
	pSdkMesh->vertexHeaders[0].NumVertices = mesh->GetNumVertices();
	pSdkMesh->vertexHeaders[0].SizeBytes = vDesc.Size;
	pSdkMesh->vertexHeaders[0].StrideBytes = D3DXGetDeclVertexSize(declaration, 0);

	pSdkMesh->subsets[0].IndexCount = pSdkMesh->indexHeaders[0].NumIndices;
	pSdkMesh->subsets[0].IndexStart = 0;
	pSdkMesh->subsets[0].VertexCount = pSdkMesh->vertexHeaders[0].NumVertices;
	pSdkMesh->subsets[0].VertexStart = 0;
	pSdkMesh->subsets[0].PrimitiveType = pSdkMesh->indexHeaders[0].IndexType;
	pSdkMesh->subsets[0].MaterialID = 0;

	pSdkMesh->meshes[0].IndexBuffer = 0;
	pSdkMesh->meshes[0].VertexBuffers[0] = 0;
	pSdkMesh->meshes[0].NumSubsets = 1;
	pSdkMesh->meshes[0].NumVertexBuffers = 1;

	pSdkMesh->meshHeader->NumIndexBuffers = 1;
	pSdkMesh->meshHeader->NumMaterials = nMaterials;
	pSdkMesh->meshHeader->NumMeshes = 1;
	pSdkMesh->meshHeader->NumTotalSubsets = 1;
	pSdkMesh->meshHeader->NumVertexBuffers = 1;

	void* bufferData;
	hr = vertexBuffer->Lock(0,0,&bufferData, NULL);
	if (FAILED(hr)) goto END;
	pSdkMesh->vertexDatas = (BYTE**)malloc((size_t)pSdkMesh->vertexHeaders[0].SizeBytes);
	memcpy(pSdkMesh->vertexDatas, bufferData, (size_t)pSdkMesh->vertexHeaders[0].SizeBytes);
	hr = vertexBuffer->Unlock();
	if (FAILED(hr)) goto END;
	hr = indexBuffer->Lock(0,0,&bufferData, NULL);
	if (FAILED(hr)) goto END;
	pSdkMesh->indexDatas = (BYTE**)malloc((size_t)pSdkMesh->indexHeaders[0].SizeBytes);
	memcpy(pSdkMesh->indexDatas, bufferData, (size_t)pSdkMesh->indexHeaders[0].SizeBytes);
	hr = indexBuffer->Unlock();
	if (FAILED(hr)) goto END;

END:
	if (indexBuffer != NULL) indexBuffer->Release();
	if (vertexBuffer != NULL) vertexBuffer->Release();

	if (FAILED(hr)) return hr;
	return S_OK;
}

HRESULT Save(SDKMESH sdkMesh, Options options)
{
	HRESULT hr = S_OK;
	BYTE* data = NULL;
	HANDLE fileHandle = NULL;
	LPWSTR path = NULL;

	if (options.outputPath != NULL)
	{
		int len = wcslen(options.outputPath) + 1;
		path = (LPWSTR)malloc(sizeof(WCHAR) * len);
		wcscpy_s(path, len, options.outputPath);
	}
	else if (options.inputPath != NULL)
	{
		int len = wcslen(options.inputPath) + 5;
		path = (LPWSTR)malloc(sizeof(WCHAR) * len);
		wcscpy_s(path, len, options.inputPath);
		LPWSTR dotLoc = wcsrchr(path, L'.');
		if (dotLoc == NULL)
		{
			hr = E_FAIL;
			goto END;
		}
		wcscpy_s(dotLoc, 7, L".model");
	}
	else
	{
		int len = wcslen(L".model") + 1;
		path = (LPWSTR)malloc(sizeof(WCHAR) * len);
		wcscpy_s(path, len, L".model");
	}

	sdkMesh.meshHeader->MeshDataOffset = sizeof(SDKMESH_HEADER);
	sdkMesh.meshHeader->SubsetDataOffset = sdkMesh.meshHeader->MeshDataOffset + sizeof(SDKMESH_MESH);
	sdkMesh.meshHeader->MaterialDataOffset = sdkMesh.meshHeader->SubsetDataOffset + sizeof(SDKMESH_SUBSET);
	sdkMesh.meshHeader->VertexStreamHeadersOffset = sdkMesh.meshHeader->MaterialDataOffset + (sizeof(SDKMESH_MATERIAL) * sdkMesh.meshHeader->NumMaterials);
	sdkMesh.meshHeader->IndexStreamHeadersOffset = sdkMesh.meshHeader->VertexStreamHeadersOffset + sizeof(SDKMESH_VERTEX_BUFFER_HEADER);
	sdkMesh.vertexHeaders[0].DataOffset = sdkMesh.meshHeader->IndexStreamHeadersOffset + sizeof(SDKMESH_INDEX_BUFFER_HEADER);
	sdkMesh.indexHeaders[0].DataOffset = sdkMesh.vertexHeaders[0].DataOffset + sdkMesh.vertexHeaders[0].SizeBytes;

	sdkMesh.meshes[0].SubsetOffset = sdkMesh.meshHeader->SubsetDataOffset;

	UINT64 totalSize = sdkMesh.indexHeaders[0].DataOffset + sdkMesh.indexHeaders[0].SizeBytes;

	data = (BYTE*)malloc(sizeof(BYTE)*(size_t)totalSize);
	memcpy(data, sdkMesh.meshHeader, sizeof(SDKMESH_HEADER));
	memcpy(data + sdkMesh.meshHeader->MeshDataOffset, sdkMesh.meshes, sizeof(SDKMESH_MESH));
	memcpy(data + sdkMesh.meshHeader->SubsetDataOffset, sdkMesh.subsets, sizeof(SDKMESH_SUBSET));
	memcpy(data + sdkMesh.meshHeader->MaterialDataOffset, sdkMesh.materials, sizeof(SDKMESH_MATERIAL) * sdkMesh.meshHeader->NumMaterials);
	memcpy(data + sdkMesh.meshHeader->VertexStreamHeadersOffset, sdkMesh.vertexHeaders, sizeof(SDKMESH_VERTEX_BUFFER_HEADER));
	memcpy(data + sdkMesh.meshHeader->IndexStreamHeadersOffset, sdkMesh.indexHeaders, sizeof(SDKMESH_INDEX_BUFFER_HEADER));
	memcpy(data+sdkMesh.vertexHeaders[0].DataOffset, sdkMesh.vertexDatas, (size_t)sdkMesh.vertexHeaders[0].SizeBytes);
	memcpy(data+sdkMesh.indexHeaders[0].DataOffset, sdkMesh.indexDatas, (size_t)sdkMesh.indexHeaders[0].SizeBytes);

	fileHandle = CreateFile(path, FILE_WRITE_DATA, NULL, NULL, CREATE_ALWAYS, FILE_FLAG_SEQUENTIAL_SCAN, NULL);

	if (fileHandle == INVALID_HANDLE_VALUE)
	{
		CloseHandle(fileHandle);
		hr = E_FAIL;
		goto END;
	}

	DWORD bytesWritten;
	WriteFile(fileHandle, data, (DWORD)totalSize, &bytesWritten, NULL);

	if (bytesWritten != totalSize)
	{
		CloseHandle(fileHandle);
		hr = E_FAIL;
		goto END;
	}

	CloseHandle(fileHandle);
	fileHandle = NULL;

END:
	if (data != NULL) free(data);
	if (fileHandle != NULL) CloseHandle(fileHandle);
	if (path != NULL) free(path);

	if (FAILED(hr)) return hr;
	return S_OK;
}
